# -*- coding: utf-8 -*-
# @Function: 清空表格存储数据库
# @Description: 清空table_storage表中的所有数据，提供安全的表格数据清理功能
# @Usage: 直接运行此脚本清空表格数据，或被其他模块导入使用
# @Dependencies: Z_config.py, Z_db_manager.py

import sys
import os
from Z_db_manager import DBManager

class TableDeleter:
    def __init__(self):
        """初始化删除器"""
        self.db_manager = DBManager()
    
    # def delete_all_tables(self, confirm=False):
    def delete_all_tables(self, confirm=True):
        """清空所有表格数据
        
        Args:
            confirm (bool): 是否确认删除，默认False需要用户确认
            
        Returns:
            bool: 删除是否成功
        """
        try:
            # 连接数据库
            self.db_manager.connect()
            
            # 如果没有确认，询问用户
            if not confirm:
                
                # 查询当前记录数
                count = self._get_record_count()
                if count == 0:
                    print("表格存储数据库中没有数据，无需清空。")
                    return True
                    
                print(f"当前表格存储数据库中共有 {count} 条表格记录")
                # response = input("确定要清空所有表格数据吗？(y/N): ").strip().lower()
                response = 'y'
                if response not in ['y', 'yes']:
                    print("操作已取消")
                    return False
            
            # 执行删除操作
            print("正在清空表格数据库...")
            deleted_count = self._execute_delete()
            print(f"成功清空 {deleted_count} 条表格记录")

            return True
            
        except Exception as e:
            print(f"清空表格数据库失败: {e}")
            return False
        finally:
            # 关闭数据库连接
            self.db_manager.close()
    
    def _get_record_count(self):
        """获取当前记录数量
        
        Returns:
            int: 记录数量
        """
        count_sql = "SELECT COUNT(*) FROM table_storage;"
        
        try:
            with self.db_manager.conn.cursor() as cursor:
                cursor.execute(count_sql)
                result = cursor.fetchone()
                return result[0] if result else 0
        except Exception as e:
            print(f"查询表格记录数量失败: {e}")
            return 0
    
    def _execute_delete(self):
        """执行删除操作
        
        Returns:
            int: 删除的记录数量
        """
        # 先获取要删除的记录数
        count = self._get_record_count()
        
        # 执行删除
        delete_sql = "DELETE FROM table_storage;"
        
        try:
            with self.db_manager.conn.cursor() as cursor:
                cursor.execute(delete_sql)
                self.db_manager.conn.commit()
                print(f"成功删除 {count} 条表格记录")
                return count
        except Exception as e:
            print(f"执行删除操作失败: {e}")
            self.db_manager.conn.rollback()
            raise
    
    def delete_by_category(self, category_id, confirm=False):
        """按类别删除表格数据
        
        Args:
            category_id (str): 要删除的类别ID
            confirm (bool): 是否确认删除
            
        Returns:
            bool: 删除是否成功
        """
        try:
            # 连接数据库
            self.db_manager.connect()
            
            # 查询该类别的记录数
            count_sql = "SELECT COUNT(*) FROM table_storage WHERE category_id = %s;"
            
            with self.db_manager.conn.cursor() as cursor:
                cursor.execute(count_sql, (category_id,))
                count = cursor.fetchone()[0]
                
                if count == 0:
                    print(f"类别 '{category_id}' 中没有表格数据，无需删除。")
                    return True
                
                if not confirm:
                    print(f"类别 '{category_id}' 中共有 {count} 条表格记录")
                    response = input(f"确定要删除类别 '{category_id}' 的所有表格数据吗？(y/N): ").strip().lower()
                    if response not in ['y', 'yes']:
                        print("操作已取消")
                        return False
                
                # 执行删除
                delete_sql = "DELETE FROM table_storage WHERE category_id = %s;"
                cursor.execute(delete_sql, (category_id,))
                self.db_manager.conn.commit()
                
                print(f"成功删除类别 '{category_id}' 中的 {count} 条表格记录")
                return True
                
        except Exception as e:
            print(f"删除类别表格数据失败: {e}")
            self.db_manager.conn.rollback()
            return False
        finally:
            self.db_manager.close()

    def delete_by_source_file(self, source_file, confirm=False):
        """按来源文件删除表格数据
        
        Args:
            source_file (str): 要删除的来源文件名
            confirm (bool): 是否确认删除
            
        Returns:
            bool: 删除是否成功
        """
        try:
            # 连接数据库
            self.db_manager.connect()
            
            # 查询该文件的记录数
            count_sql = "SELECT COUNT(*) FROM table_storage WHERE source_file = %s;"
            
            with self.db_manager.conn.cursor() as cursor:
                cursor.execute(count_sql, (source_file,))
                count = cursor.fetchone()[0]
                
                if count == 0:
                    print(f"来源文件 '{source_file}' 中没有表格数据，无需删除。")
                    return True
                
                if not confirm:
                    print(f"来源文件 '{source_file}' 中共有 {count} 条表格记录")
                    response = input(f"确定要删除来源文件 '{source_file}' 的所有表格数据吗？(y/N): ").strip().lower()
                    if response not in ['y', 'yes']:
                        print("操作已取消")
                        return False
                
                # 执行删除
                delete_sql = "DELETE FROM table_storage WHERE source_file = %s;"
                cursor.execute(delete_sql, (source_file,))
                self.db_manager.conn.commit()
                
                print(f"成功删除来源文件 '{source_file}' 中的 {count} 条表格记录")
                return True
                
        except Exception as e:
            print(f"删除来源文件表格数据失败: {e}")
            self.db_manager.conn.rollback()
            return False
        finally:
            self.db_manager.close()

def main():
    """主函数"""
    print("=== 表格存储数据清理工具 ===")
    
    deleter = TableDeleter()
    
    # 提供选项菜单
    while True:
        # print("\n请选择操作:")
        # print("1. 清空所有表格数据")
        # print("2. 按类别删除表格数据")
        # print("3. 按来源文件删除表格数据")
        # print("4. 退出")
        
        # choice = input("请输入选项 (1-4): ").strip()
        choice = '1'

        if choice == '1':
            deleter.delete_all_tables()
            break
        elif choice == '2':
            category_id = input("请输入要删除的类别ID: ").strip()
            if category_id:
                deleter.delete_by_category(category_id)
            else:
                print("类别ID不能为空")
        elif choice == '3':
            source_file = input("请输入要删除的来源文件名: ").strip()
            if source_file:
                deleter.delete_by_source_file(source_file)
            else:
                print("来源文件名不能为空")
        elif choice == '4':
            print("退出程序")
            break
        else:
            print("无效选项，请重新选择")

if __name__ == "__main__":
    main()